# -*- coding:utf-8  -*-
# @Author: ChenYangMing
# @Time: 2024-09-25

import pandas as pd

from Common.log_utils import LogUtils
from Common.time_utils import timeit, current_formatted_time
from config import TaskConfig
from data_process import TaskDatasets
from model import Model, train_model, predict
from evaluation import TaskEvaluator


log_name = f'{TaskConfig.CAPABILITY}_{TaskConfig.TASK}_{TaskConfig.TASK_MODEL}_{current_formatted_time()}.log'
logger = LogUtils.setup_logger(log_name)


@timeit
def main():
    """
    主函数，用于执行模型训练和预测的完整流程。
    该函数首先加载配置，构建数据集，然后对每个城市进行模型训练和预测。
    最后，使用预测结果和真实标签计算评估指标。
    """
    config = TaskConfig()
    logger.info('载入配置')
    dataset = TaskDatasets(logger)
    logger.info('开始构造训练集+测试集')
    dataset.create_dataset(config.DATA_PATH, config.HOLIDAYS_DATA_PATH)
    logger.info('数据构造完成')

    logger.info('根据配置创建算法实例')
    if config.TASK_MODEL not in config.MODEL_NAMES:
        logger.error("未知模型错误")
        raise ValueError("未知模型错误")
    predict_res = []
    for city in ['A', 'B', 'C']:
        model = Model(config, logger, dataset.input_dim, dataset.output_dim)
        # 加载预训练模型或开始新的训练
        if config.LOAD_MODEL_FROM and city in config.LOAD_MODEL_FROM:
            model_path = config.LOAD_MODEL_FROM.get(city)
            model = model.load(model_path)
        else:
            logger.info(f"city {city} 模型开始训练")
            model = train_model(model, dataset.datasets[city]['train'], logger)
            model.save(config.MODEL_PATH_TO.get(city))
        # 进行预测
        predict_res_city = predict(model, dataset.datasets[city]['test'], dataset.scaler)
        predict_res_city['地市'] = f'city{city}'
        predict_res.append(predict_res_city)
    # 合并所有城市的预测结果，并重置索引
    predict_res = pd.concat(predict_res, axis=0)
    predict_res.reset_index(drop=True, inplace=True)
    # 读取真实标签数据
    y_true = dataset.read_file(TaskConfig.LABEL_PATH, file_type='csv')
    # 创建评估器并计算评估指标
    evaluator = TaskEvaluator(y_true, predict_res)
    index = evaluator.evaluate()
    logger.info(f"模型指标: {index}")


if __name__ == "__main__":
    main()
